package org.example.slave.connector.mysql;

import cn.hutool.core.util.ArrayUtil;
import lombok.Data;
import lombok.EqualsAndHashCode;
import org.example.slave.channel.MyChannel;
import org.example.slave.channel.MyChannelPool;
import org.example.slave.connector.AbstractConnector;
import org.example.slave.constants.MysqlAuthAlgo;
import org.example.slave.packet.*;
import org.example.slave.util.MyChannelUtil;
import org.example.slave.util.MySQLPasswordEncrypter;
import org.example.slave.util.MysqlCodecUtil;

import java.net.InetSocketAddress;
import java.security.NoSuchAlgorithmException;
import java.util.Objects;

/**
 * @author dingrui
 * @since 2022/2/16
 * @description
 */
@Data
@EqualsAndHashCode(callSuper = true)
public class MysqlConnector extends AbstractConnector {

    // read timeout(ms)
    private static final long timeout = 5_000;

    private InetSocketAddress socketAddress;
    private String username;
    private String pwd;
    private byte charsetNumber = 33;
    private String defaultSchema;


    private MysqlConnector() {
    }

    public MysqlConnector(InetSocketAddress addr, String username, String pwd) {
        this.socketAddress = addr;
        this.username = username;
        this.pwd = pwd;
    }

    public MysqlConnector(InetSocketAddress addr, String username, String pwd, byte charsetNumber, String defaultSchema) {
        this(addr, username, pwd);
        this.charsetNumber = charsetNumber;
        this.defaultSchema = defaultSchema;
    }


    @Override
    protected boolean doConnect() {
        if (!super.connectedStatus.compareAndSet(false, true)) return false;
        try {
            super.myChannel = MyChannelPool.open(this.socketAddress);
            this.negotiate(this.myChannel);
            return true;
        } catch (Exception e) {
            super.disConnect();
            return false;
        }
    }

    // 接收mysql服务端握手初始化包->发送登录请求->确认登录结果->发送业务查询
    private void negotiate(MyChannel ch) {
        this.handshakeInitResponse(ch);
        this.authRequest(ch);
        this.authResponse(ch);
    }

    // 接收mysql服务端握手包
    private void handshakeInitResponse(MyChannel ch) {
        BasePacket basePacket = MysqlCodecUtil.readBasePacket(ch, 4, timeout);
        assert Objects.nonNull(basePacket) && Objects.nonNull(basePacket.getHeaderPacket()) && ArrayUtil.isNotEmpty(basePacket.getBody());
        byte[] body = basePacket.getBody();
        // check field_count
        byte marker = body[0];
        if (marker < 0) {
            if (marker == -1)
                throw new RuntimeException("err: " + MysqlCodecUtil.createErrPacket(body));
            else if (body[0] == -2)
                throw new RuntimeException("EOF");
            else
                throw new RuntimeException("field_count = " + body[0]);
        }
        // 客户端 -> mysql handshake
        HandshakePacket handshakePacket = new HandshakePacket(body);
        // 构建auth请求
        AuthPacket authPacket = new AuthPacket();
        authPacket.setCharsetNumber(this.charsetNumber);
        authPacket.setUsername(this.username);
        authPacket.setPassword(this.pwd);
        authPacket.setServerCapabilities(handshakePacket.getServerCapabilities());
        authPacket.setDatabaseName(this.defaultSchema);
        authPacket.setScrumbleBuff(this.joinAndCreateScrumbleBuff(handshakePacket));
        authPacket.setAuthPluginName(MysqlAuthAlgo.MYSQL_NATIVE_PASSWORD.getBytes());
        byte[] authBody = authPacket.toBytes();
        assert Objects.nonNull(authBody);
        // 登录请求
        byte lastSeqNo = basePacket.getHeaderPacket().getPacketSequenceNumber();
        MyChannelUtil.write(ch, authBody, lastSeqNo + 1);
    }

    private byte[] joinAndCreateScrumbleBuff(HandshakePacket handshakePacket) {
        byte[] dest = new byte[handshakePacket.seed.length + handshakePacket.restOfScrambleBuff.length];
        System.arraycopy(handshakePacket.seed, 0, dest, 0, handshakePacket.seed.length);
        System.arraycopy(handshakePacket.restOfScrambleBuff, 0, dest, handshakePacket.seed.length, handshakePacket.restOfScrambleBuff.length);
        return dest;
    }

    private void authRequest(MyChannel ch) {
        BasePacket basePacket = MysqlCodecUtil.readBasePacket(ch, 4, timeout);
        assert Objects.nonNull(basePacket) && Objects.nonNull(basePacket.getHeaderPacket()) && ArrayUtil.isNotEmpty(basePacket.getBody());
        byte[] body = basePacket.getBody();
        byte marker = body[0];
        if (marker == -2 || marker == 1) {
            byte[] authData = null;
            if (marker == 1) {
                AuthSwitchRequestMoreData authSwitchRequestMoreDate = MysqlCodecUtil.createAuthSwitchRequestMoreDate(body);
                authData = authSwitchRequestMoreDate.getAuthData();
            } else {
                AuthSwitchRequestPacket authSwitchRequestPacket = MysqlCodecUtil.createAuthSwitchRequestPacket(body);
                authData = authSwitchRequestPacket.getAuthData();
            }
            byte[] encryptedPassword = null;
            // XXX: 2022/2/15 只实现auth 加密算法mysql_native_password
            try {
                encryptedPassword = MySQLPasswordEncrypter.scramble411(this.getPwd().getBytes(), authData);
            } catch (NoSuchAlgorithmException e) {
                throw new RuntimeException("encrypt err: ", e);
            }
            assert Objects.nonNull(encryptedPassword);
            AuthSwitchResponsePacket authSwitchResponsePacket = new AuthSwitchResponsePacket();
            authSwitchResponsePacket.setAuthData(encryptedPassword);
            byte[] auth = authSwitchResponsePacket.toBytes();
            byte lastSeqNo = basePacket.getHeaderPacket().getPacketSequenceNumber();
            // 请求
            MyChannelUtil.write(ch, auth, lastSeqNo + 1);
        } else {
            if (marker < 0) {
                if (marker == -1)
                    throw new RuntimeException("err: " + MysqlCodecUtil.createErrPacket(body));
                else
                    throw new RuntimeException("field_count = " + body[0]);
            }
        }
    }

    private void authResponse(MyChannel ch) {
        // 确认登录成功 拿到请一个登录请求的响应 不要校验响应包的内容
        BasePacket basePacket = MysqlCodecUtil.readBasePacket(ch, 4, timeout);
        assert Objects.nonNull(basePacket) && ArrayUtil.isNotEmpty(basePacket.getBody());
        byte[] body = basePacket.getBody();
        byte marker = body[0];
        if (marker < 0) {
            if (marker == -1)
                throw new RuntimeException("err: " + MysqlCodecUtil.createErrPacket(body));
            else
                throw new RuntimeException("field_count = " + body[0]);
        }
    }
}
